#include "md5.h"
#include <arm_neon.h>
#include <iostream>
#include <cstring>

// 假设md5.h中已经定义了Byte、bit32等类型以及相关常量和宏

// 定义Neon指令优化的函数
inline uint32x4_t F_neon(uint32x4_t x, uint32x4_t y, uint32x4_t z) {
    return vandq_u32(x, y) | vandq_u32(vmvnq_u32(x), z);
}

inline uint32x4_t G_neon(uint32x4_t x, uint32x4_t y, uint32x4_t z) {
    return vandq_u32(x, z) | vandq_u32(y, vmvnq_u32(z));
}

inline uint32x4_t H_neon(uint32x4_t x, uint32x4_t y, uint32x4_t z) {
    return veorq_u32(veorq_u32(x, y), z);
}

inline uint32x4_t I_neon(uint32x4_t x, uint32x4_t y, uint32x4_t z) {
    return veorq_u32(y, vorrq_u32(x, vmvnq_u32(z)));
}

inline uint32x4_t ROTATELEFT_neon(uint32x4_t num, int n) {
    return vorrq_u32(vshlq_n_u32(num, n), vshrq_n_u32(num, 32 - n));
}

inline void FF_neon(uint32x4_t& a, uint32x4_t b, uint32x4_t c, uint32x4_t d, uint32x4_t x, int s, uint32x4_t ac) {
    a = vaddq_u32(a, vaddq_u32(F_neon(b, c, d), vaddq_u32(x, ac)));
    a = ROTATELEFT_neon(a, s);
    a = vaddq_u32(a, b);
}

inline void GG_neon(uint32x4_t& a, uint32x4_t b, uint32x4_t c, uint32x4_t d, uint32x4_t x, int s, uint32x4_t ac) {
    a = vaddq_u32(a, vaddq_u32(G_neon(b, c, d), vaddq_u32(x, ac)));
    a = ROTATELEFT_neon(a, s);
    a = vaddq_u32(a, b);
}

inline void HH_neon(uint32x4_t& a, uint32x4_t b, uint32x4_t c, uint32x4_t d, uint32x4_t x, int s, uint32x4_t ac) {
    a = vaddq_u32(a, vaddq_u32(H_neon(b, c, d), vaddq_u32(x, ac)));
    a = ROTATELEFT_neon(a, s);
    a = vaddq_u32(a, b);
}

inline void II_neon(uint32x4_t& a, uint32x4_t b, uint32x4_t c, uint32x4_t d, uint32x4_t x, int s, uint32x4_t ac) {
    a = vaddq_u32(a, vaddq_u32(I_neon(b, c, d), vaddq_u32(x, ac)));
    a = ROTATELEFT_neon(a, s);
    a = vaddq_u32(a, b);
}


Byte* StringProcessBatch(string inputs[4], int* n_byte) {
    // 假设四个输入字符串中最长的长度决定最终的填充长度
    int max_length = 0;
    for (int i = 0; i < 4; ++i) {
        if (inputs[i].length() > max_length) {
            max_length = inputs[i].length();
        }
    }

    int bitLength = max_length * 8;
    int paddingBits = bitLength % 512;
    if (paddingBits > 448) {
        paddingBits = 512 - (paddingBits - 448);
    }
    else if (paddingBits < 448) {
        paddingBits = 448 - paddingBits;
    }
    else if (paddingBits == 448) {
        paddingBits = 512;
    }
    int paddingBytes = paddingBits / 8;
    int paddedLength = max_length + paddingBytes + 8;

    Byte* paddedMessage = new Byte[4 * paddedLength];
    for (int i = 0; i < 4; ++i) {
        int length = inputs[i].length();
        Byte* blocks = (Byte*)inputs[i].c_str();
        // 复制原始消息
        memcpy(paddedMessage + i * paddedLength, blocks, length);
        // 添加填充字节
        paddedMessage[i * paddedLength + length] = 0x80;
        memset(paddedMessage + i * paddedLength + length + 1, 0, paddingBytes - 1);
        // 添加消息长度（64比特，小端格式）
        for (int j = 0; j < 8; ++j) {
            paddedMessage[i * paddedLength + length + paddingBytes + j] = ((uint64_t)length * 8 >> (j * 8)) & 0xFF;
        }
    }

    *n_byte = paddedLength;
    return paddedMessage;
}


void MD5HashBatch(string inputs[4], bit32 states[4][4]) {
    int messageLength;
    Byte* paddedMessage = StringProcessBatch(inputs, &messageLength);

    int n_blocks = messageLength / 64;

    uint32x4_t a = vdupq_n_u32(0x67452301);
    uint32x4_t b = vdupq_n_u32(0xefcdab89);
    uint32x4_t c = vdupq_n_u32(0x98badcfe);
    uint32x4_t d = vdupq_n_u32(0x10325476);

    for (int i = 0; i < n_blocks; ++i) {
        uint32x4_t x[16];
        for (int j = 0; j < 16; ++j) {
            uint32_t values[4];
            for (int k = 0; k < 4; ++k) {
                values[k] = (paddedMessage[4 * j + i * 64 + k * messageLength]) |
                    (paddedMessage[4 * j + 1 + i * 64 + k * messageLength] << 8) |
                    (paddedMessage[4 * j + 2 + i * 64 + k * messageLength] << 16) |
                    (paddedMessage[4 * j + 3 + i * 64 + k * messageLength] << 24);
            }
            x[j] = vld1q_u32(values);
        }

        uint32x4_t aa = a, bb = b, cc = c, dd = d;

        // Round 1
        FF_neon(aa, bb, cc, dd, x[0], s11, vdupq_n_u32(0xd76aa478));
        FF_neon(dd, aa, bb, cc, x[1], s12, vdupq_n_u32(0xe8c7b756));
        FF_neon(cc, dd, aa, bb, x[2], s13, vdupq_n_u32(0x242070db));
        FF_neon(bb, cc, dd, aa, x[3], s14, vdupq_n_u32(0xc1bdceee));
        FF_neon(aa, bb, cc, dd, x[4], s11, vdupq_n_u32(0xf57c0faf));
        FF_neon(dd, aa, bb, cc, x[5], s12, vdupq_n_u32(0x4787c62a));
        FF_neon(cc, dd, aa, bb, x[6], s13, vdupq_n_u32(0xa8304613));
        FF_neon(bb, cc, dd, aa, x[7], s14, vdupq_n_u32(0xfd469501));
        FF_neon(aa, bb, cc, dd, x[8], s11, vdupq_n_u32(0x698098d8));
        FF_neon(dd, aa, bb, cc, x[9], s12, vdupq_n_u32(0x8b44f7af));
        FF_neon(cc, dd, aa, bb, x[10], s13, vdupq_n_u32(0xffff5bb1));
        FF_neon(bb, cc, dd, aa, x[11], s14, vdupq_n_u32(0x895cd7be));
        FF_neon(aa, bb, cc, dd, x[12], s11, vdupq_n_u32(0x6b901122));
        FF_neon(dd, aa, bb, cc, x[13], s12, vdupq_n_u32(0xfd987193));
        FF_neon(cc, dd, aa, bb, x[14], s13, vdupq_n_u32(0xa679438e));
        FF_neon(bb, cc, dd, aa, x[15], s14, vdupq_n_u32(0x49b40821));

        // Round 2
        GG_neon(aa, bb, cc, dd, x[1], s21, vdupq_n_u32(0xf61e2562));
        GG_neon(dd, aa, bb, cc, x[6], s22, vdupq_n_u32(0xc040b340));
        GG_neon(cc, dd, aa, bb, x[11], s23, vdupq_n_u32(0x265e5a51));
        GG_neon(bb, cc, dd, aa, x[0], s24, vdupq_n_u32(0xe9b6c7aa));
        GG_neon(aa, bb, cc, dd, x[5], s21, vdupq_n_u32(0xd62f105d));
        GG_neon(dd, aa, bb, cc, x[10], s22, vdupq_n_u32(0x2441453));
        GG_neon(cc, dd, aa, bb, x[15], s23, vdupq_n_u32(0xd8a1e681));
        GG_neon(bb, cc, dd, aa, x[4], s24, vdupq_n_u32(0xe7d3fbc8));
        GG_neon(aa, bb, cc, dd, x[9], s21, vdupq_n_u32(0x21e1cde6));
        GG_neon(dd, aa, bb, cc, x[14], s22, vdupq_n_u32(0xc33707d6));
        GG_neon(cc, dd, aa, bb, x[3], s23, vdupq_n_u32(0xf4d50d87));
        GG_neon(bb, cc, dd, aa, x[8], s24, vdupq_n_u32(0x455a14ed));
        GG_neon(aa, bb, cc, dd, x[13], s21, vdupq_n_u32(0xa9e3e905));
        GG_neon(dd, aa, bb, cc, x[2], s22, vdupq_n_u32(0xfcefa3f8));
        GG_neon(cc, dd, aa, bb, x[7], s23, vdupq_n_u32(0x676f02d9));
        GG_neon(bb, cc, dd, aa, x[12], s24, vdupq_n_u32(0x8d2a4c8a));

        // Round 3
        HH_neon(aa, bb, cc, dd, x[5], s31, vdupq_n_u32(0xfffa3942));
        HH_neon(dd, aa, bb, cc, x[8], s32, vdupq_n_u32(0x8771f681));
        HH_neon(cc, dd, aa, bb, x[11], s33, vdupq_n_u32(0x6d9d6122));
        HH_neon(bb, cc, dd, aa, x[14], s34, vdupq_n_u32(0xfde5380c));
        HH_neon(aa, bb, cc, dd, x[1], s31, vdupq_n_u32(0xa4beea44));
        HH_neon(dd, aa, bb, cc, x[4], s32, vdupq_n_u32(0x4bdecfa9));
        HH_neon(cc, dd, aa, bb, x[7], s33, vdupq_n_u32(0xf6bb4b60));
        HH_neon(bb, cc, dd, aa, x[10], s34, vdupq_n_u32(0xbebfbc70));
        HH_neon(aa, bb, cc, dd, x[13], s31, vdupq_n_u32(0x289b7ec6));
        HH_neon(dd, aa, bb, cc, x[0], s32, vdupq_n_u32(0xeaa127fa));
        HH_neon(cc, dd, aa, bb, x[3], s33, vdupq_n_u32(0xd4ef3085));
        HH_neon(bb, cc, dd, aa, x[6], s34, vdupq_n_u32(0x4881d05));
        HH_neon(aa, bb, cc, dd, x[9], s31, vdupq_n_u32(0xd9d4d039));
        HH_neon(dd, aa, bb, cc, x[12], s32, vdupq_n_u32(0xe6db99e5));
        HH_neon(cc, dd, aa, bb, x[15], s33, vdupq_n_u32(0x1fa27cf8));
        HH_neon(bb, cc, dd, aa, x[2], s34, vdupq_n_u32(0xc4ac5665));

        // Round 4
        II_neon(aa, bb, cc, dd, x[0], s41, vdupq_n_u32(0xf4292244));
        II_neon(dd, aa, bb, cc, x[7], s42, vdupq_n_u32(0x432aff97));
        II_neon(cc, dd, aa, bb, x[14], s43, vdupq_n_u32(0xab9423a7));
        II_neon(bb, cc, dd, aa, x[5], s44, vdupq_n_u32(0xfc93a039));
        II_neon(aa, bb, cc, dd, x[12], s41, vdupq_n_u32(0x655b59c3));
        II_neon(dd, aa, bb, cc, x[3], s42, vdupq_n_u32(0x8f0ccc92));
        II_neon(cc, dd, aa, bb, x[10], s43, vdupq_n_u32(0xffeff47d));
        II_neon(bb, cc, dd, aa, x[1], s44, vdupq_n_u32(0x85845dd1));
        II_neon(aa, bb, cc, dd, x[8], s41, vdupq_n_u32(0x6fa87e4f));
        II_neon(dd, aa, bb, cc, x[15], s42, vdupq_n_u32(0xfe2ce6e0));
        II_neon(cc, dd, aa, bb, x[6], s43, vdupq_n_u32(0xa3014314));
        II_neon(bb, cc, dd, aa, x[13], s44, vdupq_n_u32(0x4e0811a1));
        II_neon(aa, bb, cc, dd, x[4], s41, vdupq_n_u32(0xf7537e82));
        II_neon(dd, aa, bb, cc, x[11], s42, vdupq_n_u32(0xbd3af235));
        II_neon(cc, dd, aa, bb, x[2], s43, vdupq_n_u32(0x2ad7d2bb));
        II_neon(bb, cc, dd, aa, x[9], s44, vdupq_n_u32(0xeb86d391));

        a = vaddq_u32(a, aa);
        b = vaddq_u32(b, bb);
        c = vaddq_u32(c, cc);
        d = vaddq_u32(d, dd);
    }

    uint32_t a_results[4], b_results[4], c_results[4], d_results[4];
    vst1q_u32(a_results, a);
    vst1q_u32(b_results, b);
    vst1q_u32(c_results, c);
    vst1q_u32(d_results, d);

    for (int i = 0; i < 4; ++i) {
        states[i][0] = a_results[i];
        states[i][1] = b_results[i];
        states[i][2] = c_results[i];
        states[i][3] = d_results[i];
    }

    delete[] paddedMessage;
}